import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
import scienceplots

# 使用SCI风格
plt.style.use(['science', 'no-latex'])

# ---------------------- 1. 数据准备 ----------------------
# 生成时间轴（x轴）
t = np.linspace(0, 8*np.pi, 500)  # 时间范围：包含4个周期，足够展示时域波形
# 生成频率轴（y轴）
freq = np.fft.fftfreq(len(t), t[1]-t[0])
freq_pos = freq[freq >= 0]  # 只保留正频率
# 构建时间-频率网格（3D图的x-y平面）
T, F = np.meshgrid(t, freq_pos)

# 定义两个时域信号（带相位偏移对比）
y1 = np.sin(t)                  # 原始信号：sin(x)
y2 = np.sin(t + np.pi/2)        # 相位偏移信号：sin(x + π/2)

# 傅里叶变换（获取频域数据）
y1f = np.fft.fft(y1)
y2f = np.fft.fft(y2)
# 提取正频率的振幅和相位（z轴数据）
amp1 = 2.0/len(t) * np.abs(y1f[freq >= 0])
amp2 = 2.0/len(t) * np.abs(y2f[freq >= 0])
phase1 = np.angle(y1f[freq >= 0])
phase2 = np.angle(y2f[freq >= 0])

# 构建3D图的z轴数据（振幅/相位扩展为网格维度）
AMP1 = np.tile(amp1.reshape(-1, 1), len(t))  # 振幅1的网格
AMP2 = np.tile(amp2.reshape(-1, 1), len(t))  # 振幅2的网格
PHASE1 = np.tile(phase1.reshape(-1, 1), len(t))  # 相位1的网格
PHASE2 = np.tile(phase2.reshape(-1, 1), len(t))  # 相位2的网格

# ---------------------- 2. 创建3D图（振幅+相位对比） ----------------------
fig = plt.figure(figsize=(16, 8), dpi=300)

# 子图1：3D振幅谱（对比两个信号的振幅分布）
ax1 = fig.add_subplot(121, projection='3d')
# 绘制sin(x)的振幅表面
surf1 = ax1.plot_surface(T, F, AMP1, cmap='viridis', alpha=0.7, label='sin(x)')
# 绘制sin(x+π/2)的振幅表面（沿x轴偏移，避免完全重叠）
surf2 = ax1.plot_surface(T+0.2, F, AMP2, cmap='plasma', alpha=0.7, label='sin(x+π/2)')
# 叠加时域波形投影（在y=0平面，方便关联时域和频域）
ax1.plot(t, np.zeros_like(t), y1, 'b-', linewidth=2, label='sin(x) 时域')
ax1.plot(t+0.2, np.zeros_like(t), y2, 'r--', linewidth=2, label='sin(x+π/2) 时域')

# 设置子图1参数
ax1.set_xlabel('Time', fontsize=10)
ax1.set_ylabel('Frequency (Hz)', fontsize=10)
ax1.set_zlabel('Amplitude', fontsize=10)
ax1.set_title('3D Amplitude Spectrum (Phase Shift Has No Impact)', fontsize=12)
ax1.set_ylim(0, 0.5)  # 聚焦主频率范围
ax1.legend()
# 调整视角（elev=仰角，azim=方位角）
ax1.view_init(elev=20, azim=45)

# 子图2：3D相位谱（对比两个信号的相位分布）
ax2 = fig.add_subplot(122, projection='3d')
# 绘制sin(x)的相位表面
surf3 = ax2.plot_surface(T, F, PHASE1, cmap='viridis', alpha=0.7, label='sin(x)')
# 绘制sin(x+π/2)的相位表面（沿x轴偏移）
surf4 = ax2.plot_surface(T+0.2, F, PHASE2, cmap='plasma', alpha=0.7, label='sin(x+π/2)')
# 叠加时域波形投影
ax2.plot(t, np.zeros_like(t), y1, 'b-', linewidth=2, label='sin(x) 时域')
ax2.plot(t+0.2, np.zeros_like(t), y2, 'r--', linewidth=2, label='sin(x+π/2) 时域')
# 标注相位偏移量（π/2≈1.57）
ax2.text(10, 0.16, np.pi/2, 'Phase Shift: π/2 ≈1.57', fontsize=10, color='red')

# 设置子图2参数
ax2.set_xlabel('Time', fontsize=10)
ax2.set_ylabel('Frequency (Hz)', fontsize=12)
ax2.set_zlabel('Phase (radians)', fontsize=10)
ax2.set_title('3D Phase Spectrum (Phase Shift Clearly Visible)', fontsize=12)
ax2.set_ylim(0, 0.5)  # 聚焦主频率范围
ax2.legend()
ax2.view_init(elev=20, azim=45)

# 添加颜色条（增强可读性）
cbar1 = fig.colorbar(surf1, ax=ax1, shrink=0.5, aspect=20)
cbar1.set_label('Amplitude', fontsize=10)
cbar2 = fig.colorbar(surf3, ax=ax2, shrink=0.5, aspect=20)
cbar2.set_label('Phase (radians)', fontsize=10)

plt.tight_layout()
plt.savefig('3d_phase_spectrum.png', dpi=300)
plt.show()

# 打印关键频率的相位值（验证偏移）
main_freq = 1/(2*np.pi)  # 正弦函数的主频率
idx = np.argmin(np.abs(freq_pos - main_freq))
print(f"主频率 {main_freq:.4f} Hz 处：")
print(f"sin(x) 相位：{phase1[idx]:.4f} 弧度")
print(f"sin(x+π/2) 相位：{phase2[idx]:.4f} 弧度")
print(f"相位差：{phase2[idx]-phase1[idx]:.4f} 弧度（理论值：π/2≈1.5708）")

